# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/24 15:15
# @Author  : Dell
# @File    : edit_question_rerank.py
# @Software: PyCharm
# @Desc    :问题改写（rag融合-rerank）也是业界常用的方式
import os
import json
from typing import Union
from operator import itemgetter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document

llm = ChatOpenAI(model_name='qwen-plus', base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", temperature=0.5,
                 api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37")

# --------------文档基础处理，索引和检索开始-----------------------
# 1、加载文本
loader = TextLoader(file_path="xiyouji.txt",encoding="utf-8")
blog_docs = loader.load()
# 2、分割文档
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=300,chunk_overlap=50)
splits = text_splitter.split_documents(blog_docs)
# 3、构建一个支持语义搜索的向量数据库（数据库在内存中）
embd=DashScopeEmbeddings(
    model="text-embedding-v1", dashscope_api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37"
)
vectorstore = Chroma.from_documents(documents=splits,embedding=embd)
#4、检索
# 将一个向量数据库（vectorstore）转换为一个检索器（retriever），并设置检索参数：返回最相似的5个结果。
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
# --------------文档基础处理，索引和检索结束-----------------------

#---------------进行问题改写和重排序开始---------------------------
template = """你是一名 AI 语言模型助手。你的任务是生成给定用户问题的五个不同版本，
以从向量数据库中检索相关文档。通过生成用户问题的多个视角，你
的目标是帮助用户克服基于距离的相似性搜索的一些限制。
提供这些以换行符分隔的备选问题。原始问题：{question}
"""
prompt = ChatPromptTemplate.from_template(template)
# 按提示词模版进行问题改写  ，生成5个问题
generate_queries = (
    prompt
    | llm
    | StrOutputParser()
    | (lambda x: x.split("\n"))
)
question = "根据这篇小说，你能介绍一下孙悟空的生活现状吗？"


def reciprocal_rank_fusion(results: list[list[Union[Document, dict]]], k=60):
    """ Reciprocal_rank_fusion 采用多个排名文档列表
         以及 RRF 公式中使用的可选参数 k """
    # 初始化一个字典来保存每个唯一文档的融合分数
    fused_scores = {}
    # 迭代每个排名文档列表
    for docs in results:
        # 迭代列表中的每个文档及其排名（列表中的位置）
        # enumerate(docs):将 docs列表转换为 (index, value)元组的迭代器 list(enumerate(["a", "b", "c"]))  # 输出 [(0, "a"), (1, "b"), (2, "c")]
        for rank, doc in enumerate(docs):
            # 安全序列化文档
            try:
                doc_str = doc.model_dump_json() if isinstance(doc, Document) else json.dumps(doc)
            except Exception:
                doc_str = str(doc)  # 终极兜底

            # 计算分数
            fused_scores[doc_str] = fused_scores.get(doc_str, 0) + 1 / (rank + k)
    # 根据融合分数对文档进行降序排序，得到最终的重新排序结果
    reranked_results = [(json.loads(doc), score) for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)]

    # 将重新排序的结果作为元组列表返回，每个元组包含文档及其融合分数
    return reranked_results
# 创建链，实现重排打分
retrieval_chain_rag_fusion = generate_queries | retriever.map() | reciprocal_rank_fusion
docs = retrieval_chain_rag_fusion.invoke({"question": question})
print(f"重排序的结果如下：\n")
for doc in docs:
    print(f"分数={doc[1]},对应检索结果片段：{str(doc[0])}\n")
#---------------进行问题改写和重排序开始---------------------------
#---------------构建rag链进行问题回答开始-----------------------------
# RAG
template1 = """根据此上下文回答以下问题：
{context}
问题：{question}
"""

prompt1 = ChatPromptTemplate.from_template(template1)

final_rag_chain = (
    {"context": retrieval_chain_rag_fusion,
     "question": itemgetter("question")}
    | prompt1
    | llm
    | StrOutputParser()
)
question1 ="根据这篇小说，你能介绍一下孙悟空的生活现状吗？"
f_re=final_rag_chain.invoke({"question":question1})
print(f"最后答复结果：{f_re}")
#---------------构建rag链进行问题回答结束-----------------------------